import cv2
from utils.commonUtils import get_all_file_in_dir


class ImgStyleChange(object):
    def __init__(self, model_name):
        # 加载模型
        # self.net = cv2.dnn.readNetFromTorch('./models/instance_norm/{0}.t7'.format(model_name))
        self.net = cv2.dnn.readNetFromTorch('../utils/style_change/models/instance_norm/{0}.t7'.format(model_name))
        self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)

    def img_dir_style_change(self, input_path, output_path):
        img_path_list = get_all_file_in_dir(input_path)
        for i, img_path in enumerate(img_path_list):
            # 进度
            print('style_change_trans: {0}%'.format(round((i + 1) / len(img_path_list) * 100, 2)))
            self.img_style_change(input_name=img_path, output_name=img_path.replace(input_path, output_path))

    def img_style_change(self, input_name, output_name):
        # 读取图片
        image = cv2.imread(input_name)
        (h, w) = image.shape[:2]
        blob = cv2.dnn.blobFromImage(image, 1.0, (w, h), (103.939, 116.779, 123.680), swapRB=False, crop=False)

        # 进行计算
        self.net.setInput(blob)
        out = self.net.forward()
        out = out.reshape(3, out.shape[2], out.shape[3])
        out[0] += 103.939
        out[1] += 116.779
        out[2] += 123.68
        out /= 255
        out = out.transpose(1, 2, 0)

        # # 输出图片
        # cv2.imshow('Styled image', out)
        # cv2.waitKey(0)

        # 保存图片
        out = out * 255
        cv2.imwrite(output_name, out)


if __name__ == '__main__':
    model = 'candy'
    input_path = './test_img/inp/'
    output_path = './test_img/oup/'
    isc = ImgStyleChange(model)
    isc.img_dir_style_change(input_path, output_path)
